import numpy as np
from .reward_function_base import BaseRewardFunction

class MissilePostureReward(BaseRewardFunction):
    """
    该奖励函数与导弹的姿态和速度相关
    并根据导弹的速度衰减来分配奖励
    """
    def __init__(self, config):
        super().__init__(config)
        self.previous_missile_v = None
        
    def reset(self, task, env):
        self.previous_missile_v = None
        return super().reset(task, env)

    def get_reward(self, task, env, agent_id):
        """
        根据导弹速度衰减和导弹与飞机速度向量的夹角来计算奖励值
        奖励的计算如下：
            首先,检查代理是否受到导弹的警告。如果有导弹警告,那么执行以下步骤,否则奖励为0。
            获取导弹的速度 missile_v 和飞机的速度 aircraft_v。
            如果 previous_missile_v 为 None,将 previous_missile_v 设置为当前的导弹速度。
            计算速度衰减 v_decrease,它表示上一时刻导弹速度与当前导弹速度之间的差异,根据差异的绝对值来计算。这个差异值经过了一些缩放,即除以 340(假设是单位为马赫的速度)并乘以 reward_scale。
            计算角度 angle,表示导弹速度向量与飞机速度向量之间的夹角的余弦值。
            如果 angle 小于0,奖励被设置为 angle 除以 max(v_decrease, 0) + 1。这个奖励值将受到速度衰减的影响。
            如果 angle 大于等于0,奖励被设置为 angle 乘以 max(v_decrease, 0)。这个奖励值将受到速度衰减的影响。
            最后,将奖励值添加到奖励轨迹中,并返回该奖励值。
        """
        reward = 0
        missile_sim = env.agents[agent_id].check_missile_warning()
        if missile_sim is not None:
            missile_v = missile_sim.get_velocity()
            aircraft_v = env.agents[agent_id].get_velocity()
            if self.previous_missile_v is None:
                self.previous_missile_v = missile_v
            v_decrease = (np.linalg.norm(self.previous_missile_v) - np.linalg.norm(missile_v)) / 340 * self.reward_scale
            angle = np.dot(missile_v, aircraft_v) / (np.linalg.norm(missile_v) * np.linalg.norm(aircraft_v))
            if angle < 0:
                reward = angle / (max(v_decrease, 0) + 1)
            else:
                reward = angle * max(v_decrease, 0)
        else:
            self.previous_missile_v = None
            reward = 0
        self.reward_trajectory[agent_id].append([reward])
        return reward